package com.atlin.project.aop;

import com.atlin.project.annotation.TrafficStatistics;
import com.atlin.project.config.CountingHttpServletResponseWrapper;
import com.atlin.project.model.domain.cds.ClientCode;
import com.atlin.project.model.domain.cds.ClientFile;
import com.atlin.project.utils.cds.TrafficCounterUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.util.Arrays;

import static com.atlin.project.model.enums.cds.TrafficStatisticsType.DOWNLOAD;
import static com.atlin.project.model.enums.cds.TrafficStatisticsType.UPLOAD;

@Aspect
@Component
@Slf4j
public class TrafficStatisticsAspect {

    @Resource
    private TrafficCounterUtils trafficCounterUtils;

    @Around("@annotation(trafficStatistics)")
    public Object around(ProceedingJoinPoint joinPoint, TrafficStatistics trafficStatistics) throws Throwable {
        Object result;
        long traffic = 0L;

        ClientCode clientCode = null;
        ClientFile clientFile = null;

        // 先解析参数
        for (Object arg : joinPoint.getArgs()) {
            if (arg instanceof ClientCode) {
                clientCode = (ClientCode) arg;
            }
            if (arg instanceof ClientFile) {
                clientFile = (ClientFile) arg;
            }
        }

        if (UPLOAD == trafficStatistics.type()) {
            // 上传统计
            for (Object arg : joinPoint.getArgs()) {
                if (arg instanceof MultipartFile) {
                    traffic = ((MultipartFile) arg).getSize();
                }
            }
            result = joinPoint.proceed();

            String userKey = (clientCode != null && StringUtils.isNotBlank(clientCode.getCode())) ? clientCode.getCode() : "unknown";
            log.info("用户 [{}] 上传流量统计: {} bytes", userKey, traffic);

            // 写入 Redis，可累加
            trafficCounterUtils.addUploadTraffic(userKey, traffic);

        } else if (DOWNLOAD == trafficStatistics.type()) {
            // 下载统计：包装 response
            for (Object arg : joinPoint.getArgs()) {
                if (arg instanceof HttpServletResponse) {
                    HttpServletResponse response = (HttpServletResponse) arg;
                    CountingHttpServletResponseWrapper responseWrapper = new CountingHttpServletResponseWrapper(response);
                    Object[] newArgs = Arrays.stream(joinPoint.getArgs())
                            .map(a -> a instanceof HttpServletResponse ? responseWrapper : a)
                            .toArray();

                    result = joinPoint.proceed(newArgs);

                    traffic = responseWrapper.getByteCount();
                    String userKey = (clientFile != null && StringUtils.isNotBlank(clientFile.getClientCode())) ? clientFile.getClientCode() : "unknown";

                    log.info("用户 [{}] 下载流量统计: {} bytes", userKey, traffic);

                    // 写入 Redis，可累加
                    trafficCounterUtils.addDownloadTraffic(userKey, traffic);
                    return result;
                }
            }
            result = joinPoint.proceed();

        } else {
            result = joinPoint.proceed();
        }
        return result;
    }
}
